# -*- coding: utf-8 -*-  
'''
评价函数
    这个评价函数优点玄了，貌似得有文科生的功底才知道怎么评价。。。

@author: luoyi
Created on 2021年3月18日
'''
import tensorflow as tf

import utils.conf as conf


#    seq2seq评价函数
class Seq2SeqMetric(tf.keras.metrics.Metric):
    
    
    def __init__(self,
                 name='Seq2SeqMetrics',
                 batch_size=conf.DATASET.get_batch_size(),
                 max_len=conf.TEXT.get_max_len(),
                 **kwargs):
        super(Seq2SeqMetric, self).__init__(name=name, **kwargs)
        
        self._batch_size = batch_size
        self._max_len = max_len
        
        self._acc = self.add_weight('acc', initializer='zero', dtype=tf.float64)
        pass


    def update_state(self, y_true, y_pred, sample_weight=None):
        '''
            @param y_true: Tensor (batch_size, max_len)
                                    batch_size：批量大小
                                    max_len：最大句子长度。不够的补0
            @param y_pred: Tensor (batch_size, max_len, word_dict_size)
                                    batch_size：批量大小
                                    max_len：最大句子长度。不够的补0
                                    word_dict_size：句子中的每个词应该是什么，在字典中的预测概率（很大）
        '''
        #    计算预测正确的字
        word_idx_pred = tf.math.argmax(y_pred, axis=-1)
        word_idx_pred = tf.cast(word_idx_pred, dtype=tf.int32)
        T = tf.math.count_nonzero(tf.math.equal(word_idx_pred, tf.cast(y_true, dtype=tf.int32)))
        TP = self._batch_size * self._max_len
        acc = T / TP
        self._acc.assign(acc)
        pass
    
    
    def reset_states(self):
        self._acc.assign(0.)
    def result(self):
        return self._acc
    
    pass


